# -*- coding: utf-8 -*-

import numpy as np
import ptycho
from pyE17 import io
import sys

basepath = '/data/CDI/cSAXS_sxdm_2010_02_partial_coherence'

# General parameters
verbose_level = 3
data_type = 'single'                    # Floating point precition. One of 'single' or 'double'
run_label = '0'

# Physical dimensions
energy = 6.2                            # Energy (in keV)
z = 7.046                                # Distance from object to screen 
ds = 172e-6                             # Camera pixel size

scans = ['S01712_S01728']
asize = np.array((128,128))             # [COULD BE AUTOMATED] Diffraction pattern array size 
ctr = asize/2                           # Diffraction pattern center position. If None: defaults to center of array
upsample_data = None                    # Needs to be documented...
flip_data = False                       # Matlab-related flipping problem - This option is probably obsolete now because it is taken care of in the code

# Scan parameters

# Positions are read directly from prepared data file
scan_type = None                        # ['round', 'raster', 'round_roi', 'custom' or None]

# Multiple-scan pattern.
#scan_pattern = (0,0)                    # Unique probe and object over all scans
#scan_pattern = (1,0)                   # one probe per scan, one (shared) object
scan_pattern = (0,0)                   # single (shared) probe, multiple objects 
#scan_pattern = (3,0)                   # single probe over three consecutive scans, single object
#scan_pattern = None                    # For any more complicated case
#scan_pattern_probe_indices = [...]
#scan_pattern_object_indices = [...]

Nmodes_probe = 1                        # Number of probe modes
Nmodes_object = 1                       # Number of probe modes

# Initial probe
initial_probe_type = None             # One of 'disc', 'gaussian', 'file', 'focus' (not implemented), or None
initial_probe = ptycho.load_run(basepath + '/analysis/S01693/S01693_128x128_0_ML.h5')['probe'][0,0]

# Initial object
initial_object_type = 'ones'            # One of 'ones', 'stxm' (not implemented), 'file', or None

#average_probes = True                  # Needs to be documented...
#average_probe_amp = 1e-2                # Needs to be documented...

dp_shift = False                        # If more than one scan, allow for sub-pixel shift between diffraction patterns of the scans.

probe_support = True                    # Probe support switch
probe_support_type = 'disc'             # [MAY NOT WORK PROPERLY] Either 'disc' or None 
probe_support_start = 0                 # [NOT IMPLEMENTED YET] Number of iterations before turning on (or off?) the probe support 
probe_support_area = .7                 # Area (relative to total area) of probe support

# I/O, interactions
pathdir_patt = basepath + '/analysis/%(scan)s/'
datafile_patt = '%(path)s%(scan)s_data_%(a0)03dx%(a1)03d.h5'  # datafile pattern
save_dir_patt = basepath + '/analysis/%(write_dir)s/'
dump_interval = 1000                 # Interval for dumping intermediate results
dump_patt = '%(run_name)s_dump.hd5'
plot_interval = 1                   # interval for plotting
doplot = True
last_plot = False                   # Whether to show a (blocking) last plot of the reconstruction
dump_plot = True                    # Dumping an image file of the plot
dump_plot_patt = scans[0] + '_' + run_label + '_dump.pdf'
dump_object = False                 # Whether to dump the object to disk
dump_object_interval = 1            # Interval for doing so

# Reconstruction (general)
# Not used for now
# algorithm = 'ML'
# algorithm = 'DM'
numit = 150

subpix = True                       # Sub-pixel positioning of the probe (slows down the reconstruction!)
subpix_method = 'linear'            # sub-pixel interpolation method, one of 'linear' or 'fourier'.
subpix_start = 0                   # Number of iteration to complete before turning on subpixels

# Reconstruction (DM)
fourier_relax_factor = .3          # This parameter should be used instead of pbound
pbound = None #1e-4                 # Power bound to be applied in the Fourier projection (old definition)
clip_object = False                 # Object clipping switch
clip_max = 1.0                      # Maximum amplitude if clipping is enabled
clip_min = 0.01                     # Minimum amplitude if clipping is enabled
average_start = 100                 # Number of iterations to complete before starting averaging
average_interval = 10               # Sampling interval for averaging 
probe_change_start = 5              # Number of iterations before turning on the probe refinement
DM_smooth_amplitude = .1

# Reconstruction (ML)
ML_type = 'Gauss'                   # ML model type one of 'Gauss', 'Poisson' or 'Euclid'
quad_interval = 10000               # Interval to compute real gradient instead of quadratic approximation. (relevant only for 'Gauss')
scale_precond = False                # object/probe preconditioning - seems especially useful for weak objects.
object_smooth_gradient = 0          # Object gradient smoothing preconditioner. (0: off, <0: Hann window, >0: Gaussian filter)
#scale_probe_object = 1.            # Fixed object/probe scale
#float_intens = True                # Let overall intensity float freely to compensate for eventual intensity fluctuations 
reg_del2 = True                     # Quadratic regularization 
reg_del2_amplitude = .001            # Quadratic regularization amplitude - probably in the range (.001 - .1)

###############################################################
# Prepare parameters
###############################################################

p = ptycho.prepare_params(globals())
ptycho.verbose(1,ptycho.print_summary(p))

###############################################################
# Run reconstruction
###############################################################

p = ptycho.ptycho_DM(p, numit=5)
p = ptycho.ptycho_ML(p, numit=30)
